package com.caochenlei.tree.huffman;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

//哈夫曼编码的实现
public class HuffmanCode {
    //哈夫曼树的根结点
    private Node root;
    //记录最后字节位数
    private int lastByteBits;

    public HuffmanCode(HuffmanTree huffmanTree) {
        this.root = huffmanTree.getRoot();
    }

    //构建哈夫曼编码表
    public Map<Byte, String> buildHuffmanCode() {
        Map<Byte, String> huffmanCodes = new HashMap<>();
        buildHuffmanCode(huffmanCodes, root, "");
        return huffmanCodes;
    }

    //构建哈夫曼编码表
    private void buildHuffmanCode(Map<Byte, String> huffmanCodes, Node node, String s) {
        if (node.isLeaf()) {
            huffmanCodes.put(node.data, s);
            return;
        }
        buildHuffmanCode(huffmanCodes, node.left, s + "0");
        buildHuffmanCode(huffmanCodes, node.right, s + "1");
    }

    //使用哈夫曼编码表编码得到编码后的字符串
    public String buildHuffmanCodeString(byte[] bytes, Map<Byte, String> huffmanCodes) {
        StringBuilder huffmanCodeString = new StringBuilder();
        for (byte data : bytes) {
            huffmanCodeString.append(huffmanCodes.get(data));
        }
        return huffmanCodeString.toString();
    }

    //将编码后的字符串转化为相对应的字节数组
    public byte[] buildHuffmanCodeBytes(String huffmanCodeString) {
        lastByteBits = huffmanCodeString.length() % 8 == 0 ? 8 : huffmanCodeString.length() % 8;
        int len = (huffmanCodeString.length() + 7) / 8;
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;
        String element;
        for (int i = 0; i < huffmanCodeString.length(); i += 8) {
            if (i + 8 > huffmanCodeString.length()) {
                element = huffmanCodeString.substring(i);
            } else {
                element = huffmanCodeString.substring(i, i + 8);
            }
            huffmanCodeBytes[index++] = (byte) Integer.parseInt(element, 2);
        }
        return huffmanCodeBytes;
    }

    //压缩
    public byte[] encode(byte[] sourceBytes, Map<Byte, String> huffmanCodes) {
        return buildHuffmanCodeBytes(buildHuffmanCodeString(sourceBytes, huffmanCodes));
    }

    //解压
    public byte[] decode(byte[] targetBytes) {
        return rebuildHuffmanCodeStringToSourceBytes(rebuildHuffmanCodeTargetBytesToString(targetBytes));
    }

    //反向将编码后的字节数组转化为编码字符串
    public String rebuildHuffmanCodeTargetBytesToString(byte[] huffmanCodeBytes) {
        String binaryString;
        StringBuilder huffmanCodeString = new StringBuilder();
        for (int i = 0; i < huffmanCodeBytes.length; i++) {
            int element = huffmanCodeBytes[i];
            if (i == huffmanCodeBytes.length - 1) {
                if (element > 0) {//如果byte是正数，需要补充满8位，负数则不需要
                    element |= 256;
                    binaryString = Integer.toBinaryString(element);
                    binaryString = binaryString.substring(binaryString.length() - lastByteBits);
                } else {
                    binaryString = Integer.toBinaryString(element);
                }
            } else {
                if (element > 0) {//如果byte是正数，需要补充满8位，负数则不需要
                    element |= 256;
                    binaryString = Integer.toBinaryString(element);
                    binaryString = binaryString.substring(binaryString.length() - 8);
                } else {
                    binaryString = Integer.toBinaryString(element);
                    binaryString = binaryString.substring(binaryString.length() - 8);
                }
            }
            huffmanCodeString.append(binaryString);
        }
        return huffmanCodeString.toString();
    }

    //反向将编码字符串转化为源字符串字节数组
    public byte[] rebuildHuffmanCodeStringToSourceBytes(String huffmanCodeString) {
        List<Byte> datas = new ArrayList<>();
        Node node = this.root;
        for (int i = 0; i < huffmanCodeString.length(); i++) {
            char flag = huffmanCodeString.charAt(i);
            if (flag == 48) {       //如果是0，则代表data在左子树中
                node = node.left;
            }
            if (flag == 49) {       //如果是1，则代表data在右子树中
                node = node.right;
            }
            if (node.isLeaf()) {    //如果是叶子，则找到data，重置node
                datas.add(node.data);
                node = this.root;
            }
        }
        byte[] huffmanCodeBytes = new byte[datas.size()];
        for (int i = 0; i < datas.size(); i++) {
            huffmanCodeBytes[i] = datas.get(i);
        }
        return huffmanCodeBytes;
    }
}
